# -*- coding: utf-8 -*-
import time
import os,sys
import pwd,stat
from otp import *

# 日志
def logprint(*msg):
    pass
    print(*msg)

# 读取配置，格式每行一组k=v
def get_config(path):
    config = {}
    with open(path, 'r') as f:
        while True:
            line = f.readline()
            if not line:
                break
            if line.startswith('#'):
                continue
            idx = line.find("=")
            if idx >= 0:
                a = line[0:idx]
                b = line[idx+1:]
                config[a.strip()] = b.strip()
    return config

# 检查totp口令
def check(username, token):
    if not username or not token:
        return 1

    #### 获取用户
    uu = pwd.getpwnam(username)
    if not uu:
        logprint('用户没找到')
        return 1

    #### 获取路径
    if not uu.pw_dir:
        logprint('用户目录没找到')
        return 1

    path = uu.pw_dir + '/.pam_mytotp.conf'
    if not os.path.isfile(path):
        logprint('totp配置文件不存在')
        return 1

    ### 获取totp配置文件信息，要求文件权限为0600，文件所属为当前认证用户
    fst = os.stat(path)
    mode = fst.st_mode

    if not stat.S_ISREG(mode):
        logprint('totp配置文件不是普通文件')
        return 1

    if fst.st_size <= 0 or fst.st_size > 1024 * 10:
        logprint('totp配置文件为空或文件大小超过10KB')
        return 1

    if stat.S_IMODE(mode) != (0o600):
        logprint('totp配置文件权限不是0600')
        return 1

    uid = fst.st_uid
    if fst.st_uid != uu.pw_uid:
        logprint('totp配置文件所有者不是当前认证用户')
        return 1

    #### 获取配置
    logprint('totp配置读取: ', path)
    config = get_config(path)
    key = config.get('key')
    key_format = config.get('key_format')
    window = config.get('window')
    interval = config.get('interval')

    if not key:
        logprint('otp参数项key为空')
        return 1
    if not key_format:
        key_format = "string"
    if not window:
        window = 3
    if not interval:
        interval = 30

    #### 验证totp
    t = TOTP(None, int(window), int(interval))
    t.setSecret(key, key_format)

    if t.verify(token) != True:
        logprint('totp检测失败')
        return 1
    else:
        logprint('totp检查成功')
        return 0

def main(argv):
    #time.sleep(3)
    if len(argv) < 3:
        logprint('参数个数错误')
        exit(1)
    username = argv[1]
    token = argv[2]
    ret = check(username, token)
    logprint('检测结果:', ret)
    exit(ret)

if __name__ == "__main__":
    main(sys.argv)
